# 环节3：文本向量化（改用阿里云通义千问 text-embedding-v4）
# 导入对接通义千问的必要库（修复LangChain导入路径）
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.document_loaders import PyPDFLoader  # PDF加载仍用新版路径
from langchain.text_splitter import RecursiveCharacterTextSplitter  # 改为旧版导入路径！！！


def initialize_tongyi_embeddings(api_key):
    """
    初始化阿里云通义千问 text-embedding-v4 向量化工具
    参数:
        api_key: 你的阿里云 DashScope API Key（从控制台获取）
    返回:
        embeddings: 初始化好的通义千问向量化工具对象
    """
    try:
        # 初始化通义千问嵌入模型，指定使用 text-embedding-v4
        embeddings = DashScopeEmbeddings(
            model="text-embedding-v4",  # 固定模型名称，不可修改
            dashscope_api_key=api_key  # 传入你的API密钥
        )
        print(f"✅ 成功初始化通义千问 text-embedding-v4 模型（远程API调用）")
        print(f"ℹ️  模型特性：生成1536维向量，适合中文文本语义理解")
        return embeddings
    except Exception as e:
        # 捕获API密钥错误、网络错误等常见问题
        error_msg = f"❌ 通义千问模型初始化失败：{str(e)}\n"
        error_msg += "请检查以下问题：\n"
        error_msg += "1. API Key 是否正确（从阿里云DashScope控制台获取）\n"
        error_msg += "2. 阿里云账号是否已实名认证（未实名无法调用API）\n"
        error_msg += "3. 网络是否能访问阿里云API（需避免代理拦截）"
        print(error_msg)
        raise  # 抛出异常，避免后续代码报错


def vectorize_documents_with_tongyi(embeddings, documents):
    """
    调用通义千问 API，将文档片段转成向量
    参数:
        embeddings: 初始化好的通义千问向量化工具
        documents: 分割后的文档片段列表（Document对象，来自环节2）
    返回:
        vectors: 每个文档片段对应的1536维向量列表
        vector_metadata: 向量对应的元数据（页码、片段长度，方便后续溯源）
    """
    # 1. 提取文档片段的文字内容和元数据
    texts = [doc.page_content for doc in documents]
    metadata_list = [
        {
            "page": doc.metadata.get("page", "未知"),  # 页码
            "text_length": len(doc.page_content)  # 片段字符长度
        }
        for doc in documents
    ]

    # 2. 调用通义千问 API 生成向量（批量处理，提高效率）
    print(f"\n📡 开始调用通义千问 API 处理 {len(texts)} 个文本片段...")
    try:
        # 由于通义千问API限制每次最多处理10个文本，需要分批处理
        vectors = []
        batch_size = 10  # 通义千问API限制每次最多处理10个文本
        
        # 分批处理所有文本
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            print(f"正在处理批次 {i//batch_size + 1}/{(len(texts)-1)//batch_size + 1}，包含 {len(batch_texts)} 个文本...")
            batch_vectors = embeddings.embed_documents(batch_texts)
            vectors.extend(batch_vectors)
    except Exception as e:
        error_msg = f"❌ API调用失败：{str(e)}\n"
        error_msg += "常见原因：\n"
        error_msg += "1. API Key 余额不足（阿里云通义千问有免费额度，超额需充值）\n"
        error_msg += "2. 文本片段过长（单片段建议不超过5000字符，避免API拒绝）\n"
        error_msg += "3. 网络波动（可重试1-2次）"
        print(error_msg)
        raise

    # 3. 验证向量格式（确保是1536维，避免API返回异常）
    if len(vectors) > 0 and len(vectors[0]) != 1024:
        raise ValueError(f"❌ 通义千问返回向量维度异常，应为1024维，实际为{len(vectors[0])}维")

    print(f"✅ API调用完成！共生成 {len(vectors)} 个1024维向量")
    return vectors, metadata_list


# 主程序：测试完整流程（加载PDF→分割→通义千问向量化）
if __name__ == "__main__":
    # --------------------------
    # 关键配置（请替换为你的信息）
    # --------------------------
    PDF_FILE_PATH = "平安保险用户手册.pdf"  # 你的PDF文件路径（确保存在）
    TONGYI_API_KEY = "sk-6434b554122c4b049ceb805d703f695b"  # 替换为你的阿里云API Key

    # 1. 第一步：加载PDF文档（复用修复后的加载逻辑）
    try:
        loader = PyPDFLoader(PDF_FILE_PATH)
        document_pages = loader.load()
        print(f"\n📄 成功加载PDF文件：{PDF_FILE_PATH}")
        print(f"   - 总页数：{len(document_pages)} 页")
    except FileNotFoundError:
        print(f"❌ 找不到PDF文件：{PDF_FILE_PATH}，请检查路径是否正确！")
        exit()  # 终止程序，避免后续报错

    # 2. 第二步：分割文档（与之前逻辑一致，确保片段大小适合API）
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,  # 单片段500字符（远低于API 5000字符上限）
        chunk_overlap=100,  # 片段重叠100字符（避免语义断裂）
        length_function=len  # 按字符数计算长度
    )
    split_docs = text_splitter.split_documents(document_pages)
    print(f"\n✂️  文档分割完成！")
    print(f"   - 原始页数：{len(document_pages)} 页")
    print(f"   - 分割后片段数：{len(split_docs)} 个")
    print(f"   - 单片段最大字符数：{max(len(doc.page_content) for doc in split_docs)}")

    # 3. 第三步：初始化通义千问向量化工具（关键步骤）
    try:
        tongyi_embeddings = initialize_tongyi_embeddings(TONGYI_API_KEY)
    except Exception as e:
        print(f"\n❌ 向量化工具初始化失败，程序终止：{str(e)}")
        exit()

    # 4. 第四步：调用API生成向量（测试时先用前5个片段，避免消耗过多API额度）
    SAMPLE_SIZE = 5  # 测试用：只处理前5个片段（正式使用可改为 len(split_docs)）
    sample_docs = split_docs[:SAMPLE_SIZE]
    print(f"\n📤 开始处理 {SAMPLE_SIZE} 个测试片段...")

    try:
        sample_vectors, sample_metadata = vectorize_documents_with_tongyi(
            embeddings=tongyi_embeddings,
            documents=sample_docs
        )
    except Exception as e:
        print(f"\n❌ 向量生成失败，程序终止：{str(e)}")
        exit()

    # 5. 第五步：展示结果（验证向量是否正确生成）
    print("\n----- 通义千问向量化结果预览 -----")
    for i, (doc, vec, meta) in enumerate(zip(sample_docs, sample_vectors, sample_metadata)):
        print(f"\n【片段 {i + 1}】")
        print(f"📄 来源页码：第 {meta['page'] + 1} 页（注：API返回页码从0开始）")
        print(f"📝 文本预览：{doc.page_content[:100]}..."  # 显示前100字符
              f"（总长度：{meta['text_length']} 字符）")
        print(f"🔢 向量信息：1536维向量，前10个数值：{vec[:10]}...")  # 简化输出（避免过长）